# -*- coding: utf-8 -*-
# @Time    : 2024/9/12 13:32
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : gMOT_4beam.py
# @Software: PyCharm

import matplotlib.pyplot as plt
import numpy
import pandas
from mpl_toolkits.mplot3d.axes3d import Axes3D

from mot.atom import Rb87_D2
from mot.average_scattering_force_model import MOT
from mot.field.base import RectGrid3D
from mot.field.laser import NonDiffractionLaserBeam, LaserInfo, BeamPolarization, set_frequency
from mot.field.magnetic import BGenerators
from mot.field.magnetic import CurrentCarryingLoop
from mot.gmot.grating import Sector, GratingSection, LaserDiffractedByGratingSection
from myutils.geometrical import Plane
import matplotlib
color_table = list(matplotlib.colors.TABLEAU_COLORS.keys())

plt.ion()
grid = RectGrid3D(numpy.array((*numpy.linspace(-10e-3, -5.5e-3, 10),
                               *numpy.linspace(-5e-3, 5e-3, 25),
                               *numpy.linspace(5.5e-3, 10e-3, 10)
                               )),
                  numpy.array((*numpy.linspace(-10e-3, -5.5e-3, 10),
                               *numpy.linspace(-5e-3, 5e-3, 26),
                               *numpy.linspace(5.5e-3, 10e-3, 10)
                               )),
                  numpy.array((*numpy.linspace(0, 8e-3, 25),
                               *numpy.linspace(9e-3, 15e-3, 15)
                               )
                              ),
                  [
                      # "Bx", "By", "Bz"
                  ]
                  )
# grid = RectGrid3D(numpy.array((*numpy.linspace(-10e-3, 12e-3, 35),
#                                )),
#                   numpy.array((*numpy.linspace(-10e-3, 12e-3, 36),
#                                )),
#                   numpy.array((*numpy.linspace(0, 13e-3, 40),
#                                )
#                               ),
#                   [
#                       # "Bx", "By", "Bz"
#                   ]
#                   )
grating_sections = [GratingSection(
    Sector(numpy.array([0, 0]), 10e-3, 1e-3, numpy.deg2rad(-60 + i * 120), numpy.deg2rad(60 + i * 120)).shape,
    Plane(numpy.array([0, 0, grid.zs.min()]), numpy.array([0, 0, 1]))) for i in range(3)
]
# plt.figure()
# for gs in grating_sections:
#     plt.plot(*gs.shape.exterior.xy)
# plt.gca().set_aspect('equal')
aperture = 30e-3 / 2
linedata = Rb87_D2()
detuning_Hz = -1.61 * linedata.Gamma  # -1.72 * linedata.Gamma
intensity_of_incident_laser = (17.3 * 0 + 16.8) * 1e-3 / (1e-2) ** 2
__polarizations = [+1, -1]
incident_laser = NonDiffractionLaserBeam(
    set_frequency(LaserInfo(numpy.array([0, 0, -1]), BeamPolarization(__polarizations[0]), "incident", detuning_Hz),
                  linedata.f_0 + detuning_Hz),
    Plane(numpy.array([0, 0, grid.zs.max()]), numpy.array([0, 0, 1]), ),
    lambda x, y, z: intensity_of_incident_laser if x ** 2 + y ** 2 < aperture ** 2 else 0.,
)
reflectivity = 36.3e-2  # 40.9e-2
phi = numpy.deg2rad(42.5 * 0 + 46.3)
lasers_diffracted_by_grating_section = [
    LaserDiffractedByGratingSection(
        grating_sections[i], incident_laser,
        set_frequency(LaserInfo(numpy.array(  # [-1, 0, 2]
            [*(1 * numpy.array([numpy.cos(numpy.deg2rad(180 + 120 * i)),
                                numpy.sin(numpy.deg2rad(180 + 120 * i))])),
             numpy.tan(phi)]), BeamPolarization(__polarizations[1]),
            "grating%d" % i, detuning_Hz), linedata.f_0 + detuning_Hz), reflectivity / numpy.sin(phi)
    ) for i, gs in
    enumerate(grating_sections)]
I_coil =3.71e3# 229.71598951823216 * 16.1
z0 = 4.45e-3  # 5.3e-3
R_coil = 8.48e-2
coil1 = CurrentCarryingLoop(I_coil, R_coil, numpy.array([0, 0, -100e-3]), numpy.array([0, 0, 1]))
coil2 = CurrentCarryingLoop(I_coil, R_coil, numpy.array([0, 0, 2 * z0 - coil1.center[2]]), numpy.array([0, 0, -1]))
mot_initializer = MOT(grid, BGenerators([coil1, coil2]), [incident_laser, *lasers_diffracted_by_grating_section],
                      linedata)
mot_initializer.cache()
mm = 1e-3
Gs = 1e-4


plt.figure(figsize=(4,3), constrained_layout = True)
zs__ = numpy.linspace(10e-3, -5e-3, 100)#mot_initializer.grid. zs
eps = 1e-10
Bz_on_axis =BGenerators([coil1, coil2]).B(eps,eps,zs__)[:,2]
plt.plot(-zs__  / mm  ,
  Bz_on_axis       /Gs,)
plt.axvline(0,  ls = ':', c = color_table[1],#label = "grating"
            )
plt.axvline(-6.98,  ls = ':', c = color_table[2],
            )
plt.xlabel("z / mm")
plt.ylabel("$B_z$ / Gs")
B_df = pandas.DataFrame({'z':-zs__, 'Bz':Bz_on_axis})
import sys
sys.path.append(r"F:\changeworld\HPMCalc")
from simulation.task_manager.simulator import  df_to_gdf
df_to_gdf(B_df, r"E:\CSTprojects\gMOT_CAES\pp\GPT/Bz_on_axis.gdf")

X, Z = numpy.meshgrid(grid.xs, grid.zs,
                      # numpy.linspace(grid.xs.min(), grid.xs.max(), 40),
                      # numpy.linspace(grid.zs.min(), grid.zs.max(), 41),  # indexing='ij'
                      )
X_for_streamplot, Z_for_streamplot = numpy.meshgrid(
    numpy.linspace(grid.xs.min(), grid.xs.max(), 40),
    numpy.linspace(grid.zs.min(), grid.zs.max(), 41),  # indexing='ij'
)
plt.figure()
plt.contourf(X / mm, Z / mm, grid.get(X, 0, Z, mot_initializer.laser_field_names.sumI))
plt.streamplot(X_for_streamplot / mm, Z_for_streamplot / mm,
               *grid.get(X_for_streamplot, 0, Z_for_streamplot, ["Bx", "Bz"]).transpose((2, 0, 1)))
plt.scatter(X / mm, Z / mm, s=0.1)
plt.gca().set_aspect('equal')
# F = mot_initializer.scattering_force_from_laser(
#     grid.X, grid.Y, grid.Z,
#     numpy.zeros(grid.X.shape), numpy.zeros(grid.X.shape), numpy.zeros(grid.X.shape)
# )
F_for_streamplot_XZ = mot_initializer.scattering_force_from_laser(
    X_for_streamplot, 0 * numpy.ones(X_for_streamplot.shape), Z_for_streamplot,
    numpy.zeros(X_for_streamplot.shape), numpy.zeros(X_for_streamplot.shape), numpy.zeros(X_for_streamplot.shape)
)
plt.figure()
plt.contourf(X / mm, Z / mm, grid.get(X, 0, Z, mot_initializer.laser_field_names.sumI), )
plt.streamplot(X_for_streamplot / mm, Z_for_streamplot / mm, F_for_streamplot_XZ[:, :, 0],
               F_for_streamplot_XZ[:, :, 2], )
plt.gca().set_aspect('equal')
z_ref = 5.08e-3
_X, _Y = numpy.meshgrid(grid.xs, grid.ys,
                        # numpy.linspace(grid.xs.min(), grid.xs.max(), 40),
                        # numpy.linspace(grid.zs.min(), grid.zs.max(), 41),  # indexing='ij'
                        )
_X_for_streamplot, _Y_for_streamplot = numpy.meshgrid(
    numpy.linspace(grid.xs.min(), grid.xs.max(), 40),
    numpy.linspace(grid.ys.min(), grid.ys.max(), 41),  # indexing='ij'
)
mm = 1e-3
F_for_streamplot_XY = mot_initializer.scattering_force_from_laser(
    _X_for_streamplot, _Y_for_streamplot, z_ref * numpy.ones(_X_for_streamplot.shape),
    numpy.zeros(_X_for_streamplot.shape), numpy.zeros(_X_for_streamplot.shape), numpy.zeros(_X_for_streamplot.shape)
)
plt.figure()
plt.contourf(_X / mm, _Y / mm, grid.get(_X, _Y, z_ref, mot_initializer.laser_field_names.sumI), )
plt.streamplot(_X_for_streamplot / mm, _Y_for_streamplot / mm, F_for_streamplot_XY[:, :, 0],
               F_for_streamplot_XY[:, :, 1],  # start_points=sttpts
               )
plt.gca().set_aspect('equal')
xs = numpy.linspace(grid.xs.min(), grid.xs.max(),
                    # -2.5e-3, 2.5e-3,
                    100)
zs = numpy.linspace(grid.zs.min(), grid.zs.max(),
                    100)
Bx_ = grid.get(xs, 0, z0, "Bx")
Bz_ = grid.get(0, 0, zs, 'Bz')
Gs = 1e-4  # 1高斯
cm = 1e-2
# plt.figure()
# plt.plot(xs / cm, Bx_ / Gs)
# plt.plot(zs / cm, Bz_ / Gs)
# plt.xlabel('x / cm')
# plt.ylabel('Bx / Gs')
from _logging import logger

Gx_data = numpy.diff(Bx_) / numpy.diff(xs)  # G = d(Bx)/dx, 其中d为偏微分符号
Gz_data = numpy.diff(Bz_) / numpy.diff(zs)
Gx = lambda x: numpy.interp(x, xs[1:], Gx_data)
Gz = lambda z: numpy.interp(z, zs[1:], Gz_data)
logger.info("梯度dB/dx = %.2f T/m (Bq = %.2f G/cm)" % (Gx(0), Gz(z0) / (Gs / cm)))
logger.info("四极场中心位置zoff= %.2f mm" % (z0 / mm))

ts_dumped = numpy.linspace(0, 1e-3, 200)
from scipy.integrate import solve_ivp


def dY_dt(t, Y_: numpy.ndarray):
    """

    :param t:
    :param Y: shape (6*N_particles, ), 其中每6位分别为(x, y, z, vx, vy, vz)
    :return:
    """
    Y = Y_.reshape((-1, 6))
    ret = numpy.zeros(Y.shape)
    ret[:, :3] = Y[:, -3:]
    ret[:, -3:] = mot_initializer.scattering_force_from_laser(
        *[Y[..., i] for i in range(6)]) / mot_initializer.transition_line_data.mass_kg
    return ret.reshape((-1,))


# trajdata = solve_ivp(dY_dt, (ts_dumped[0], ts_dumped[-1]), numpy.array([
#     [0, 1e-3, 1e-3, 1, 1, 1],
#     [1e-3, 1e-3, 5e-3, 1, -1, 1],
# ]).T, t_eval=ts_dumped,vectorized  = True)
N_0 = int(1e4)
position_0 = numpy.array([grid.xs.min(), grid.ys.min(), grid.zs.min()]) + (
        numpy.array([grid.xs.max(), grid.ys.max(), grid.zs.max()]) -
        numpy.array([grid.xs.min(), grid.ys.min(), grid.zs.min()])) * numpy.random.random(
    (N_0, 3))  # grid.all_points().reshape((-1, 3))
i_particles = numpy.arange(len(position_0))[numpy.random.random(len(position_0)) > 0.9]
position_0 = position_0[i_particles]
v_0 = numpy.zeros(position_0.shape)
Y0 = numpy.array((*position_0.transpose((1, 0)), *v_0.transpose((1, 0)))).transpose((1, 0))
from mot.particle_initializer import MaxwellParticleInitializer

min_and_max = lambda xs: numpy.array((xs.min(), xs.max()))

Y0 = MaxwellParticleInitializer(
    N_0,
    numpy.array([min_and_max(grid.xs),
                 min_and_max(grid.ys),
                 min_and_max(grid.zs),
                 ]), 0, 20,2, mot_initializer.transition_line_data.mass_kg).generate().values

trajdata = solve_ivp(dY_dt, (ts_dumped[0], ts_dumped[-1]), y0=Y0.reshape((-1,)), t_eval=ts_dumped, vectorized=False)
N_particles = int(len(trajdata.y[:, 0]) / 6)

fig = plt.figure()
ax: Axes3D = fig.add_subplot(projection='3d')
color = "red"
i_particles = numpy.arange(N_particles)[numpy.random.random(N_particles) > 0 * 0.99]
MOT_center = numpy.array([0, 0, 8e-3])
last_pos = trajdata.y.reshape((N_particles, 6, -1))[:, :3, -1]
last_vel = trajdata.y.reshape((N_particles, 6, -1))[:, -3:, -1]
indexes_of_trapped_particles = (numpy.linalg.norm(last_pos - MOT_center, axis=1) < 3e-3) & (
        numpy.linalg.norm(last_vel, axis=1) < 10)
for i in numpy.where(indexes_of_trapped_particles)[0]:
    ax.plot3D(*(trajdata.y.reshape((N_particles, 6, -1))[i, :3, :] / mm), c=color, lw=0.1)
    ax.scatter(*(trajdata.y.reshape((N_particles, 6, -1))[i, :3, -1] / mm), c=color)
ax.scatter(*last_pos[indexes_of_trapped_particles].T / mm, s=1)
ax.set_xlabel("x / mm")
ax.set_ylabel("y / mm")
ax.set_zlabel("z / mm")
ax.set_xlim(*min_and_max(grid.xs / mm))
ax.set_ylim(*min_and_max(grid.ys / mm))
ax.set_zlim(*min_and_max(grid.zs / mm))

capture_rate = len(last_pos[indexes_of_trapped_particles]) / N_particles
logger.info("俘获效率：%.2f%%" % (100 * capture_rate))
logger.info("被俘获粒子的最终位置分布范围：\nX[%.2fmm, %.2fmm], Y[%.2fmm, %.2fmm], Z[%.2fmm, %.2fmm]" % (
    *(min_and_max(last_pos[indexes_of_trapped_particles, 0]) / mm),
    *(min_and_max(last_pos[indexes_of_trapped_particles, 1]) / mm),
    *(min_and_max(last_pos[indexes_of_trapped_particles, 2]) / mm),))
import scipy.constants as C

logger.info("被俘获粒子的平均温度：%.2e K" % (
        mot_initializer.transition_line_data.mass_kg * numpy.linalg.norm(last_vel[indexes_of_trapped_particles],
                                                                         axis=1).mean() ** 2 / (3 * C.k)))
plt.figure()
default_size = 5
plt.scatter(Y0[:, 0] / mm, Y0[:, 3], label="start", #s=default_size
            )
plt.scatter(last_pos[:, 0] / mm, last_vel[:, 0], label="end",# s=default_size
            )
# plt.xlabel("x / mm")
# plt.ylabel("vx / mm")
# plt.legend()
#
# plt.figure()
plt.scatter(Y0[indexes_of_trapped_particles, 0] / mm, Y0[indexes_of_trapped_particles, 3], label="captured, start",
            #s=default_size
            )
plt.scatter(last_pos[indexes_of_trapped_particles, 0] / mm, last_vel[indexes_of_trapped_particles, 0],
            label="captured, end",# s=default_size
            )
plt.xlabel("x / mm")
plt.ylabel("vx / (m/s)")
plt.legend()
